Sourced from jaxtyping's releases.

jaxtyping v0.2.24

New features

• Error messages will now include useful shape information for debugging. (!!!) This closes the venerable #6, which is is one of the oldest feature requests for jaxtyping. This is enabled by using the following syntax, instead of the old double-decorator syntax:

  from jaxtyping import jaxtyped
  from beartype/typeguard import beartype/typechecked as typechecker
  @​jaxtyped(typechecker=typechecker)  # passing as keyword argument is important
  def foo(...):
  ...

  and moreover this is what install_import_hook now does. As an example of this done, consider this buggy code:

  import jax.numpy as jnp
  from jaxtyping import Array, Float, jaxtyped
  from beartype import beartype
  @jaxtyped(typechecker=beartype)
  def f(x: Float[Array, "foo bar"], y: Float[Array, "foo"]):
  ...
  f(jnp.zeros((3, 4)), jnp.zeros(5))

  will now produce the error message

  jaxtyping.TypeCheckError: Type-check error whilst checking the parameters of f.
  The problem arose whilst typechecking argument 'y'.
  Called with arguments: {'x': f32[3,4], 'y': f32[5]}
  Parameter annotations: (x: Float[Array, 'foo bar'], y: Float[Array, 'foo']).
  The current values for each jaxtyping axis annotation are as follows.
  foo=3
  bar=4
  

  Hurrah! I'm really glad to have this important quality-of-life improvement in. (#6, #138)

• Added support for the following:

  def make_zeros(size: int) -> Float[Array, "{size}"]:
      return jnp.zeros(size)

  in which axis names enclosed in {...} are evaluated as f-strings using the value of the argument of the function. This closes the long-standing feature request #93. (#93, #140) (Heads-up @​MilesCranmer!)
• Added support for declaring PyTree structures, which like array shapes must match across all arguments. For example

  def f(x: PyTree[int, "T"], y: PyTree[float, "T"])

  demands that x and y be PyTrees with the same jax.tree_util.tree_structure as each other. (#135)
• Added support for treepath-dependent sizes using ?. This makes it possible for the value of a dimension to vary across its position within a pytree, but must still be consistent with its value in other pytrees of the same structure. Such annotations look like PyTree[Float[Array, "?foo"], "T"]. Together with the previous point, this means that you can now declare that two pytrees must have the exact same structure and array shapes as each other: use PyTree[Float[Array, "?*shape"], "T"] as the annotation for both. (#136)
• Added jaxtyping.Real, which admits any float, signed integer, or unsigned integer. (But not bools or complexes.) (#128)
• If JAX is installed, then jaxtyping.DTypeLike is now available (it is just a forwarding on of jax.typing.DTypeLike). (#129)

... (truncated)

• 8e47c90 version bump
• 2059789 Fix install_import_hook(..., None)
• 850f4e7 Handle beartype doing isinstance(None, hint)
• 0a76c9c Error message improvements
• 5fbd671 Added support for 'self' in dataclass attribute annotations; switched from ar...
• 80a9956 Removed unused elements from unions, e.g. Float[ArrayLike, ...] will no longe...
• baffbef Doc fix
• 127eae5 array shapes+dtypes no longer checked as part of pytree flattening. This avoi...
• 58600d3 Fixed new-style PRNG keys.
• 7925e27 Symbolic expressions now support delayed binding to arguments. Fixes #93.
• Additional commits viewable in compare view

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)